package zutils

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"github.com/bytedance/sonic"
	"strings"
)

type JwtHeader struct {
	Algo string `json:"alg"`
	Type string `json:"typ"`
}

var DefaultHeader = JwtHeader{
	Algo: "HS256",
	Type: "JWT",
}

type JwtPayLoad struct {
	ID          string         `json:"jti"`
	Issue       string         `json:"iss"`
	Audience    string         `json:"aud"`
	Subject     string         `json:"sub"`
	IssueAt     int64          `json:"iat"`
	NotBefore   int64          `json:"nbf"`
	Expireation int64          `json:"exp"`
	UserDefined map[string]any `json:"ud"`
}

func GenJWT(header JwtHeader, payload JwtPayLoad, secret string) (string, error) {
	var part1, part2, signature string
	if bs1, err := sonic.Marshal(header); err != nil {
		return "", err
	} else {
		part1 = base64.RawURLEncoding.EncodeToString(bs1)
	}

	if bs2, err := sonic.Marshal(payload); err != nil {
		return "", err
	} else {
		part2 = base64.RawURLEncoding.EncodeToString(bs2)
	}

	hash := hmac.New(sha256.New, []byte(secret))
	hash.Write([]byte(part1 + "." + part2))

	signature = base64.RawURLEncoding.EncodeToString(hash.Sum(nil))

	return part1 + "." + part2 + "." + signature, nil
}

func VerifyJwt(token string, secret string) (*JwtHeader, *JwtPayLoad, error){
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return nil, nil, fmt.Errorf("token是 %d 部分", len(parts))
	}
	hash := hmac.New(sha256.New, []byte(secret))
	hash.Write([]byte(parts[0] + "." + parts[1]))
	signature := base64.RawURLEncoding.EncodeToString(hash.Sum(nil))
	if signature != parts[2] {
		return nil, nil, fmt.Errorf("前面验证失败！！！")
	}

	var part1, part2 []byte
	var err error
	if part1, err = base64.RawURLEncoding.DecodeString(parts[0]); err != nil {
		return nil, nil, fmt.Errorf("header Base64反解失败！！！")
	}

	if part2, err = base64.RawURLEncoding.DecodeString(parts[1]); err != nil {
		return nil, nil, fmt.Errorf("payload Base64反解失败！！！")
	}

	var header JwtHeader
	var payload JwtPayLoad
	if err = sonic.Unmarshal(part1, &header); err != nil {
		return nil, nil, fmt.Errorf("header json反解失败！！！")
	}
	if err = sonic.Unmarshal(part2, &payload); err != nil {
		return nil, nil, fmt.Errorf("payload json反解失败！！！")
	}

	return &header, &payload, nil

}